from fastmcp import FastMCP
import anyio
import random

# 创建一个天气服务器实例
weather_mcp = FastMCP(
    name="WeatherAssistant",
    instructions="""
        This server provides weather information.
        Call get_weather() to get current weather data for a location.
        Call get_forecast() to get weather forecast for a location.
    """,
)

@weather_mcp.tool
async def get_weather(city: str, country: str = "CN") -> dict:
    """获取指定城市的当前天气信息"""
    # 模拟网络请求或API调用
    def fetch_weather_data(city: str, country: str) -> dict:
        # 这里可以替换为真实的天气API调用
        temperature = random.randint(-10, 40)
        humidity = random.randint(30, 90)
        conditions = random.choice(["晴天", "多云", "雨天", "雪天", "阴天"])
        
        return {
            "city": city,
            "country": country,
            "temperature": temperature,
            "humidity": humidity,
            "conditions": conditions,
            "description": f"{city}当前天气为{conditions}，温度{temperature}°C，湿度{humidity}%"
        }
    
    # 使用 anyio.to_thread.run_sync 避免阻塞事件循环
    return await anyio.to_thread.run_sync(fetch_weather_data, city, country)

@weather_mcp.tool
async def get_forecast(city: str, days: int = 3, country: str = "CN") -> dict:
    """获取指定城市未来几天的天气预报"""
    def fetch_forecast_data(city: str, days: int, country: str) -> dict:
        forecast = []
        conditions_options = ["晴天", "多云", "雨天", "雪天", "阴天"]
        
        for i in range(days):
            day_data = {
                "day": f"第{i+1}天",
                "high": random.randint(0, 40),
                "low": random.randint(-10, 20),
                "conditions": random.choice(conditions_options)
            }
            forecast.append(day_data)
            
        return {
            "city": city,
            "country": country,
            "forecast": forecast
        }
    
    return await anyio.to_thread.run_sync(fetch_forecast_data, city, days, country)

# 添加一个简单的工具用于测试
@weather_mcp.tool
def get_supported_cities() -> list:
    """获取支持查询天气的城市列表"""
    return ["北京", "上海", "广州", "深圳", "杭州", "成都", "武汉", "西安"]

if __name__ == "__main__":
    # 运行天气服务器
    weather_mcp.run(transport='sse')